"""
Load-balancing broker

Clients and workers are shown here in-process.

Author: Brandon Carpenter (hashstat) <brandon(dot)carpenter(at)pnnl(dot)gov>
"""
#!/usr/bin/python
#-*- coding=utf-8 -*-


from __future__ import print_function

import multiprocessing

import zmq
import json
from enum import Enum
import time

class TaskType(Enum):
    A = 1
    B = 2

class WorkerStatus(Enum):
    Ready = 1
    Working = 2

class FrontCMDType(Enum):
    CS_TASK_REQ = 1
    SC_TASK_RSP = 2


def client_task(task_type):
    """Basic request-reply client using REQ socket."""
    socket = zmq.Context().socket(zmq.DEALER)
    socket.identity = u"Client-{0}".format(TaskType(task_type).name).encode("ascii")
    socket.connect("tcp://localhost:50000")

    print(socket.identity.decode('ascii'), "start")
    # protocal : json or  you can use other
    req = {'task_type': task_type, 'other':'1'}
    task_idx = 1
    for i in xrange(3):
        req['task'] = 'task_{0}_{1}'.format(TaskType(task_type).name, task_idx)
        socket.send_string(json.dumps(req))
        task_idx += 1
        # print(socket.identity.decode('ascii'), "send", req)
        print(socket.identity.decode('ascii'), "product", req['task'])

    while True:
        reply = socket.recv()
        result = json.loads(reply)
        # print(socket.identity, '--------', result)
        print("{0}: recv from server, {1} {2}".format(socket.identity.decode("ascii"), result['task'], (result['status'] == 1  and ['finish'] or ['working'])[0]))


def worker_task(task_type):
    """Worker task, using a REQ socket to do load-balancing."""
    socket = zmq.Context().socket(zmq.DEALER)
    socket.identity = u"Worker-{}".format(TaskType(task_type).name).encode("ascii")
    socket.connect("tcp://localhost:50001")

    # Tell broker we're ready for work
    status = WorkerStatus.Ready
    rsp = {'task_type': task_type, 'status': status.value}
    # socket.send_string(json.dumps(rsp))
    socket.send_string(json.dumps(rsp))
    # print(socket.identity.decode('ascii'), "send", rsp)
    print(socket.identity.decode('ascii'), "ready")

    while True:
        req = socket.recv()
        req = json.loads(req)
        # print("{}: recv  {}".format(socket.identity.decode("ascii"), req))
        if status == WorkerStatus.Working:
            rsp['status'] = status.value
            socket.send_string(json.dumps(rsp))
            # print(socket.identity.decode('ascii'), "send", rsp)
            continue

        status = WorkerStatus.Working
        rsp['status'] = status.value
        rsp['task'] = req['task']
        rsp['idx'] = req['idx']
        socket.send_string(json.dumps(rsp))
        # print(socket.identity.decode('ascii'), "send", rsp)
        print(socket.identity.decode('ascii'), "working", rsp['task'])
        time.sleep(1)
        status = WorkerStatus.Ready
        rsp['status'] = status.value
        rsp['task'] = req['task']
        rsp['client'] = req['client']
        print(socket.identity.decode('ascii'), "finish", rsp['task'])
        socket.send_string(json.dumps(rsp))
        # print(socket.identity.decode('ascii'), "send", rsp)


def main():
    """Load balancer main loop."""
    # Prepare context and sockets
    context = zmq.Context.instance()
    frontend = context.socket(zmq.ROUTER)
    frontend.bind("tcp://*:50000")
    backend = context.socket(zmq.ROUTER)
    backend.bind("tcp://*:50001")


    # Initialize main loop state
    workers = {}
    tasks = {}
    task_idx = 1
    poller = zmq.Poller()
    # Only poll for requests from backend until workers are available
    poller.register(backend, zmq.POLLIN)
    poller.register(frontend, zmq.POLLIN)

    # Start background tasks
    def start(task, *args):
        process = multiprocessing.Process(target=task, args=args)
        process.daemon = True
        process.start()
    for i in (TaskType):
        start(client_task, i.value)
    for i in (TaskType):
        start(worker_task, i.value)

    while True:
        sockets = dict(poller.poll())

        if backend in sockets:
            # Handle worker activity on the backend
            worker, req_str = backend.recv_multipart()
            req = json.loads(req_str)
            # print("proxy recv from", worker, req)
            workers[req['task_type']] = {'worker': worker, 'status': req['status']}
            if (req.get('task', None) and req.get('client', None) and req['status'] == WorkerStatus.Ready.value):
                rsp = json.dumps(req)
                frontend.send_multipart([req['client'].encode('ascii'), rsp])

            if (req.get('task', None) and req.get('idx', None) and req['status'] == WorkerStatus.Working.value):
                del tasks[req['idx']]

        if frontend in sockets:
            # Get next client request, route to last-used worker
            client, req_str = frontend.recv_multipart()
            # print("proxy recv from", client, req_str)
            req = json.loads(req_str)
            req['client'] = client
            req['idx'] = task_idx
            tasks[task_idx] = req
            task_idx = task_idx + 1

        for idx in tasks:
            worker = workers.get(tasks[idx]['task_type'], None)
            if not worker or worker['status'] != WorkerStatus.Ready.value:
                continue

            rsp = json.dumps(tasks[idx])
            backend.send_multipart([worker['worker'], rsp])
            # print("proxy send to", worker['worker'], json.dumps(tasks[idx]))
            del workers[tasks[idx]['task_type']]

    # Clean up
    backend.close()
    frontend.close()
    context.term()


if __name__ == "__main__":
    main()
